import tkinter as tk
import random

# 地图大小
w = 10
h = 10
# 食物位置
food_x = 1
food_y = 1
# 蚂蚁数量
size = 2
# 蚁群位置
ant_x = 5
ant_y = 5
# 当前蚂蚁位置数据集
ants = []
# 路径集合{x_y:信息素权重}
path = {}

# 初始化窗口
root = tk.Tk()
root.title("蜂群算法")

# 创建画布
canvas = tk.Canvas(root, width=w * 10, height=h * 10)
canvas.pack()

# 绘制食物
food = canvas.create_oval(food_x * 10, food_y * 10, (food_x + 1) * 10, (food_y + 1) * 10, fill="orange")

# 初始化蚂蚁
for i in range(size):
    # 0:蜜蜂名称，1:X轴坐标，2:Y轴坐标，3:距离，4:状态（0:随机移动、1:找到路径、2：发现食物回家、3：未发现食物回家），5:路径数组
    ants.append([i, ant_x, ant_y, 0, 0, []])
    # 显示蜜蜂的初始位置
    ant_color = f"#{random.randint(0, 0xFFFFFF):06x}"
    ant = canvas.create_oval(ant_x * 10, ant_y * 10, (ant_x + 1) * 10, (ant_y + 1) * 10, fill=ant_color)

# 最大递归次数
def timer():
    for a in ants:
        root.after(1000, ant, a)
    root.after(2000, timer)

# 单只蚂蚁的活动
def ant(a):
    # 默认随机移动
    to = -1
    # 随机移动
    if a[4] == 0:
        to = random.randint(0, 3)
    # 找到路径
    if a[4] == 1:
        max_pheromone = 0
        # 依次判断周围食物的信息素：上、右、下、左
        for i in range(4):
            xy = new_x_y(a, i)
            new_x = xy[0]
            new_y = xy[1]
            key = f"{new_x}_{new_y}"
            # 已知最强的食物信息素
            if key in path and path[key] > max_pheromone:
                max_pheromone = path[key]
                to = i
    # 更新数据
    if a[4] == 0 or a[4] == 1:
        xy = new_x_y(a, to)
        new_x = xy[0]
        new_y = xy[1]
        name = f"{new_x}_{new_y}"
        # 移动到新位置
        a[1] = new_x
        a[2] = new_y
        # 追加距离
        a[3] += 1
        # 是否有食物的信息素
        if name in path:
            a[4] = 1
        # 找到食物
        if food_x == new_x and food_y == new_y:
            a[4] = 2
            # 食物最大信息素强度
            path[name] = 999
        # 追加路径
        a[5].append(name)
        # 显示动画
        canvas.coords(ants[a[0]], new_x * 10, new_y * 10, (new_x + 1) * 10, (new_y + 1) * 10)
    # 回家，直接回家
    if a[4] == 2:
        ps = a[5]
        if ps:
            node = ps.pop()  # 删除最近一次的路径
            x_y = node.split('_')
            # 后退
            a[1] = int(x_y[0])
            a[2] = int(x_y[1])
            # 留下找到食物的信息素
            if node in path:
                # 累计增加权重
                path[node] += 1
            else:
                path[node] = 1
            # 显示动画
            canvas.coords(ants[a[0]], a[1] * 10, a[2] * 10, (a[1] + 1) * 10, (a[2] + 1) * 10)
        else:
            # 实际工程中可以适当减弱家门口的信息素
            p_name = f"{ant_x}_{ant_y}"
            path[p_name] = 0
            # 回到家里并重新开始
            a[1] = ant_x
            a[2] = ant_y
            a[3] = 0
            a[4] = 0
            print(f"蚂蚁 {a[0]} 已经回到了家里！")

# 机器人合理移动方向
def new_x_y(robot, to, len=1):
    new_x = robot[1]
    new_y = robot[2]
    # 最大边界值
    max_w = w - 1
    max_h = h - 1
    # 顺时针：上、右、下、左
    if to == 0:  # 上
        new_y = new_y - len if new_y - len >= 0 else 0
    elif to == 1:  # 右
        new_x = new_x + len if new_x + len <= max_w else max_w
    elif to == 2:  # 下
        new_y = new_y + len if new_y + len <= max_h else max_h
    elif to == 3:  # 左
        new_x = new_x - len if new_x - len >= 0 else 0
    # 返回最终的坐标点[X,Y]
    return [new_x, new_y]

# 开始运行
timer()

root.mainloop()